Custom WebSocket client

Overview

The SilentShardWebsocketClient class is a protocol-aware WebsocketClient designed specifically for MPC (Multi-Party Computation) operations. It serves as a thin wrapper around the WebsocketClient.

1. SilentShardWebsocketClient implements WebsocketClient

• Responsibility:
  • Raw WebSocket connection management
  • Message queueing and asynchronous I/O
  • Binary/text data transmission
• Customization Use Cases:
  • Add Action specific logic for example perform some logic on keygen connect
  • Replace WebSocket library (e.g., Socket.IO, different WebSocket impl)
  • Modify low-level message queueing logic

Example implementing custom WebsocketClient

App.tsx

import Foundation

open class SilentShardWebsocketClient: WebsocketClient {
    private let websocketConfig: WebsocketConfig
    private var webSocketTask: URLSessionWebSocketTask?
    private var session: URLSession? = nil
    private let logger = Logger(
        subsystem: "com.silencelaboratories.silentshard", category: "WebSocket")

    public init(websocketConfig: WebsocketConfig) {
        self.websocketConfig = websocketConfig
        self.session = URLSession(
            configuration: .default
        )
    }

    open func connect(websocketAction: WebsocketAction, params: String?)
        async throws
    {
        var urlString =
            websocketConfig.finalWebsocketBaseUrl
            + websocketAction.finalWebsocketActionUrl
        if let params = params {
            urlString += "/\(params)"
        }

        guard let url = URL(string: urlString) else {
            throw NSError(domain: "Invalid URL", code: -1)
        }

        logger.info("Connecting to \(urlString)")

        var request = URLRequest(url: url)

        request.setValue(
            "*/*",
            forHTTPHeaderField: "Accept")
        request.setValue(
            "UTF-8",
            forHTTPHeaderField: "Accept-Charset")
        request.setValue(
            websocketAction.protocol,
            forHTTPHeaderField: "Sec-WebSocket-Protocol")

        webSocketTask = session!.webSocketTask(with: request)
        webSocketTask?.resume()

        if let token = websocketConfig.authenticationToken {
            logger.info("Sending auth token")
            try await send(text: token)
            logger.info("Auth token sent")
        }
    }

    open func send(data: Data) async throws {
        guard let webSocketTask = webSocketTask else {
            logger.error("Session is null, cannot send bytes")
            throw NSError(domain: "WebSocket not connected", code: -1)
        }

        logger.info("Sending bytes: \(data.count)")
        try await webSocketTask.send(.data(data))
        logger.info("Bytes sent")
    }

    open func send(text: String) async throws {
        guard let webSocketTask = webSocketTask else {
            logger.error("Session is null, cannot send text")
            throw NSError(domain: "WebSocket not connected", code: -1)
        }

        logger.info("Sending text: \(text)")
        try await webSocketTask.send(.string(text))
        logger.info("Text sent")
    }

    open func read() async throws -> Data {
        guard let webSocketTask = webSocketTask else {
            logger.error("Session is null, cannot read")
            throw NSError(domain: "WebSocket not connected", code: -1)
        }

        logger.info("Reading from WebSocket")
        let message = try await webSocketTask.receive()

        switch message {
        case .data(let data):
            logger.info("Received bytes: \(data.count)")
            return data

        case .string(let text):
            logger.info("Received text: \(text)")
            return text.data(using: .utf8) ?? Data()

        @unknown default:
            return Data()
        }
    }

    open func disconnect() async throws {
        guard let webSocketTask = webSocketTask else {
            logger.info("Session is null, no disconnection needed")
            return
        }

        logger.info("Initiating disconnection from WebSocket")
        logger.info("Sending close frame")

        // Send our close signal
        webSocketTask
            .cancel(with: .normalClosure, reason: "fgsdfgsf".data(using: .utf8))

        logger.info("\(webSocketTask.closeCode)")

    }

    open func sendPostRequest(
        websocketAction: WebsocketAction, authHeader: String?, body: Any
    ) async throws -> Data {
        let urlString =
            websocketConfig.finalPostRequestBaseUrl
            + websocketAction.finalPostRequestActionUrl

        guard let url = URL(string: urlString) else {
            throw NSError(domain: "Invalid URL", code: -1)
        }

        logger.info("Sending POST request to \(urlString)")

        var request = URLRequest(url: url)
        request.httpMethod = "POST"

        if let authHeader = authHeader {
            logger.info("Adding auth header")
            request.setValue(
                "Bearer \(authHeader)", forHTTPHeaderField: "Authorization")
        }

        switch body {
        case let stringBody as String:
            request.setValue(
                "application/json", forHTTPHeaderField: "Content-Type")
            request.httpBody = stringBody.data(using: .utf8)
        case let dataBody as Data:
            request.setValue(
                "application/octet-stream", forHTTPHeaderField: "Content-Type")
            request.httpBody = dataBody
        default:
            throw NSError(domain: "Invalid body type", code: -1)
        }

        let (data, _) = try await session!.data(for: request)
        return data
    }
}

// Logger implementation for iOS
private struct Logger {
    let subsystem: String
    let category: String

    func info(_ message: String) {
        NSLog("[\(subsystem)][\(category)] INFO: \(message)")
    }

    func error(_ message: String) {
        NSLog("[\(subsystem)][\(category)] ERROR: \(message)")
    }
}